#!/usr/bin/python3
import os
import cv2
import numpy as np
import struct

face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/video_crop_face_list.attr.txt'
video_label_list_file = r'/rootfs/media/kasim/Data1/data/VideoGenderCropFace/video_label_list.txt'
file_list_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/faceid_person_image_gender.txt'

# face_attr_list_file = r'/rootfs/media/yery/Kaso/data/VideoGenderCropFace/video_crop_face_list.attr.txt'
# video_label_list_file = r'/rootfs/media/yery/Kaso/data/VideoGenderCropFace/video_label_list.txt'
# file_list_path = r'/rootfs/media/yery/Kaso/data/VideoGenderCropPerson/faceid_person_image_gender.txt'


# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
FILTER_QUALITY = 3
FILTER_QUALITY_THR = 0.25
QUALITY_THR = 0.3

MIN_PITCH = -20.0
MAX_PITCH = 20.0
MIN_YAW = -20.0
MAX_YAW = 20.0
MIN_ROLL = -30.0
MAX_ROLL = 30.0

ANGLE_SCORE = 0.1
ANGLE_DIFF = 10

AGE2_SCORE = 0.1
AGE_DIFF = 8

TOPK = 5


UNDER_TWO = 0
THREE_TO_FIVE = 1
SIX_TO_TWELVE = 2
THIRTEEN_TO_EIGHTEEN = 3
NINETEEN_TO_TWENTY_THREE = 4
TWENTY_FOUR_TO_TWENTY_NINE = 5
THIRTY_TO_THIRTY_FIVE = 6
THIRTY_SIX_TO_FORTY_FIVE = 7
FORTY_SIX_TO_SIXTY = 8
ABOVE_SIXTY = 9
UNCERTAIN = 10


def cal_angle(bin_angle, k=5):
    # bin_angle = np.array(bin_angle, dtype=np.float32)
    bin_angle_index0 = np.argmax(bin_angle, 1)
    idxes = np.arange(0, bin_angle.shape[0])
    bin_angle_prob0 = bin_angle[idxes, bin_angle_index0]

    end_idx = bin_angle.shape[1] - 1

    bin_angle_probs = []
    bin_angle_indexes = []
    c = 1 + k // 2
    for offset in range(1, c):
        bin_angle_index1 = bin_angle_index0 - offset
        bin_angle_index2 = bin_angle_index0 + offset

        start_mask = bin_angle_index1 < 0
        end_mask = bin_angle_index2 > end_idx

        bin_angle_index1[start_mask] = 0
        bin_angle_index2[end_mask] = 0

        bin_angle_prob1 = bin_angle[idxes, bin_angle_index1]
        bin_angle_prob2 = bin_angle[idxes, bin_angle_index2]

        bin_angle_prob1[start_mask] = 0
        bin_angle_prob2[end_mask] = 0

        bin_angle_indexes.append(bin_angle_index1)
        bin_angle_indexes.append(bin_angle_index2)
        bin_angle_probs.append(bin_angle_prob1)
        bin_angle_probs.append(bin_angle_prob2)

    bin_angle_indexes.append(bin_angle_index0)
    bin_angle_probs.append(bin_angle_prob0)

    bin_angle_prob = np.stack(bin_angle_probs, 1)
    bin_angle_prob_norm = np.sum(bin_angle_prob, 1, keepdims=True)
    bin_angle_prob = bin_angle_prob / bin_angle_prob_norm

    bin_angle_index = np.stack(bin_angle_indexes, 1)
    bin_angle = 3 * bin_angle_index - 99

    bin_angle = bin_angle_prob * bin_angle
    # bin_angle = np.sum(bin_angle, 1, keepdims=True)
    bin_angle = np.sum(bin_angle, 1)

    angle_score = bin_angle_prob0

    return bin_angle, angle_score


def cal_age2(bin_age2, k=5):
    # bin_age2 = np.array(bin_age2, dtype=np.float32)
    bin_age2_index0 = np.argmax(bin_age2, 1)
    idxes = np.arange(0, bin_age2.shape[0])
    bin_age2_prob0 = bin_age2[idxes, bin_age2_index0]

    end_idx = bin_age2.shape[1] - 1

    bin_age2_probs = []
    bin_age2_indexes = []
    c = 1 + k // 2
    for offset in range(1, c):
        bin_age2_index1 = bin_age2_index0 - offset
        bin_age2_index2 = bin_age2_index0 + offset

        start_mask = bin_age2_index1 < 0
        end_mask = bin_age2_index2 > end_idx

        bin_age2_index1[start_mask] = 0
        bin_age2_index2[end_mask] = 0

        bin_age2_prob1 = bin_age2[idxes, bin_age2_index1]
        bin_age2_prob2 = bin_age2[idxes, bin_age2_index2]

        bin_age2_prob1[start_mask] = 0
        bin_age2_prob2[end_mask] = 0

        bin_age2_indexes.append(bin_age2_index1)
        bin_age2_indexes.append(bin_age2_index2)
        bin_age2_probs.append(bin_age2_prob1)
        bin_age2_probs.append(bin_age2_prob2)

    bin_age2_indexes.append(bin_age2_index0)
    bin_age2_probs.append(bin_age2_prob0)

    bin_age2_prob = np.stack(bin_age2_probs, 1)
    bin_age2_prob_norm = np.sum(bin_age2_prob, 1, keepdims=True)
    bin_age2_prob = bin_age2_prob / bin_age2_prob_norm

    bin_age2_index = np.stack(bin_age2_indexes, 1)
    bin_age2 = 3 * bin_age2_index

    bin_age2 = bin_age2_prob * bin_age2
    # bin_age2 = np.sum(bin_age2, 1, keepdims=True)
    bin_age2 = np.sum(bin_age2, 1)

    age2_score = bin_age2_prob0

    return bin_age2, age2_score


def age_to_group(age):
    if age < 3:
        return UNDER_TWO
    if age < 6:
        return THREE_TO_FIVE
    if age < 13:
        return SIX_TO_TWELVE
    if age < 19:
        return THIRTEEN_TO_EIGHTEEN
    if age < 24:
        return NINETEEN_TO_TWENTY_THREE
    if age < 30:
        return TWENTY_FOUR_TO_TWENTY_NINE
    if age < 36:
        return THIRTY_TO_THIRTY_FIVE
    if age < 46:
        return THIRTY_SIX_TO_FORTY_FIVE
    if age < 61:
        return FORTY_SIX_TO_SIXTY
    return ABOVE_SIXTY


def main():
    video_images = {}
    image_list = []
    with open(face_attr_list_file, 'r') as file:
        file_count = 0
        for line in file.readlines():
            # gd/zs/nqjd/dxjyjy/150100414a54443452064d2742c93600/video/20190817_064924/00007_000.jpg
            file_path = line.strip()
            image_list.append(file_path)
            video = os.path.dirname(file_path)
            images = video_images.get(video, None)
            if images is None:
                video_images[video] = [file_path]
            else:
                images.append(file_path)
            file_count += 1
            # if file_count > 10000:
            #     break

    _base_path = os.path.splitext(face_attr_list_file)[0]
    title_file_path = _base_path + '.title.txt'
    with open(title_file_path, 'r') as file:
        title_str = file.readline()
        attr_item_count_str = file.readline()
        titles = title_str.strip().split(';')
        attr_item_counts = attr_item_count_str.strip().split(';')
        attr_item_counts = list(map(int, attr_item_counts))
        total_attr_item_count = sum(attr_item_counts)
        attr_count = len(titles)

    attr_item_offsets = [0]
    offset = 0
    for attr_item_count in attr_item_counts:
        offset += attr_item_count
        attr_item_offsets.append(offset)
    attrs_size = total_attr_item_count * 4
    data_file_path = _base_path + '.dat'

    image_attributes_dict = {}
    with open(data_file_path, 'rb') as file:
        for image_name in image_list:
            dat = file.read(attrs_size)
            datas = struct.unpack('{}f'.format(total_attr_item_count), dat)

            attribute = {}
            for i, title in enumerate(titles):
                start_offset = attr_item_offsets[i]
                end_offset = attr_item_offsets[i+1]
                attribute[title] = datas[start_offset:end_offset]

            image_attributes_dict[image_name] = attribute

    video_set = set()
    with open(video_label_list_file, 'r') as file:
        for line in file.readlines():
            video_name = line.split()[0].strip()
            video_set.add(video_name)

    video_ages = {}
    for video, images in video_images.items():
        if video not in video_set:
            continue
        file_names = []
        qualitys = []
        pitchs = []
        yaws = []
        rolls = []
        bin_pitchs = []
        bin_yaws = []
        bin_rolls = []
        age1s = []
        age2s = []
        for image in images:
            attributes = image_attributes_dict[image]
            quality = attributes['quality']
            pitch = attributes['pitch'][0]
            yaw = attributes['yaw'][0]
            roll = attributes['roll'][0]
            bin_pitch = attributes['bin_pitch']
            bin_yaw = attributes['bin_yaw']
            bin_roll = attributes['bin_roll']
            age1 = attributes['age1'][0]
            age2 = attributes['age2']
            file_names.append(image)
            qualitys.append(quality)
            pitchs.append(pitch)
            yaws.append(yaw)
            rolls.append(roll)
            bin_pitchs.append(bin_pitch)
            bin_yaws.append(bin_yaw)
            bin_rolls.append(bin_roll)
            age1s.append(age1)
            age2s.append(age2)

        qualitys = np.array(qualitys, dtype=np.float32)
        pitchs = np.array(pitchs, dtype=np.float32)
        yaws = np.array(yaws, dtype=np.float32)
        rolls = np.array(rolls, dtype=np.float32)
        bin_pitchs = np.array(bin_pitchs, dtype=np.float32)
        bin_yaws = np.array(bin_yaws, dtype=np.float32)
        bin_rolls = np.array(bin_rolls, dtype=np.float32)
        age1s = np.array(age1s, dtype=np.float32)
        age2s = np.array(age2s, dtype=np.float32)

        # max_idx = np.argmax(qualitys, axis=1)
        # quality_mask = (max_idx == FILTER_QUALITY) & (qualitys[max_idx] >= FILTER_QUALITY_THR)

        qualitys = 0.3333333333 * qualitys[:, 2] + 0.66666666 * qualitys[:, 1] + qualitys[:, 0]
        quality_mask = qualitys >= QUALITY_THR

        bin_pitchs, bin_pitch_scores = cal_angle(bin_pitchs)
        bin_rolls, bin_roll_scores = cal_angle(bin_rolls)
        bin_yaws, bin_yaw_scores = cal_angle(bin_yaws)

        pitch_diffs = np.abs(bin_pitchs - pitchs)
        roll_diffs = np.abs(bin_rolls - rolls)
        yaw_diffs = np.abs(bin_yaws - yaws)

        score_mask = (bin_pitch_scores > ANGLE_SCORE) & (bin_roll_scores > ANGLE_SCORE) & (bin_yaw_scores > ANGLE_SCORE)
        diff_mask = (pitch_diffs < ANGLE_DIFF) & (roll_diffs < ANGLE_DIFF) & (yaw_diffs < ANGLE_DIFF)

        pitch_mask = (MIN_PITCH <= pitchs) & (pitchs <= MAX_PITCH)
        roll_mask = (MIN_ROLL <= rolls) & (rolls <= MAX_ROLL)
        yaw_mask = (MIN_YAW <= yaws) & (yaws <= MAX_YAW)
        angle_mask = pitch_mask & roll_mask & yaw_mask

        age2s, age2_scores = cal_age2(age2s)
        age_diff = np.abs(age1s - age2s)
        age_mask = (age_diff <= AGE_DIFF) & (age2_scores >= AGE2_SCORE)

        mask = score_mask & diff_mask & quality_mask & angle_mask & age_mask
        if mask.any():
            age1s = age1s[mask]
            age2_scores = age2_scores[mask]
            video_ages[video] = [age1s, age2_scores]

    faceid_video_genders = {}
    with open(file_list_path, 'r') as file:
        for line in file.readlines():
            # 0 0 gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190912_154308/00001_000.jpg 1
            lines = line.strip().split()
            faceid = '{} {}'.format(lines[0], lines[1])
            file_path = lines[2]
            gender = lines[3]
            video_name = os.path.dirname(file_path)
            video_genders = faceid_video_genders.get(faceid, None)
            if video_genders is None:
                faceid_video_genders[faceid] = {video_name: gender}
            else:
                video_genders[video_name] = gender

    age_face_count = 0
    age_file_count = 0
    faceid_video_gender_age_agegroups = {}
    for faceid, video_genders in faceid_video_genders.items():
        age1s = []
        age2_scores = []
        face_gender = -1
        for video, gender in video_genders.items():
            face_gender = gender
            ages = video_ages.get(video, None)
            if ages is not None:
                age1s.append(ages[0])
                age2_scores.append(ages[1])
        ages_count = len(age1s)
        if ages_count > 0:
            if ages_count > 1:
                age1s = np.concatenate(age1s, axis=0)
                age2_scores = np.concatenate(age2_scores, axis=0)
            else:
                age1s = age1s[0]
                age2_scores = age2_scores[0]
            if age1s.shape[0] > TOPK:
                score_indexes = np.argsort(-age2_scores)
                score_indexes = score_indexes[0:TOPK]
                age1s = age1s[score_indexes, ...]
                age2_scores = age2_scores[score_indexes, ...]
            weights = age2_scores * age2_scores

            w = weights
            a = age1s
            sum_w = np.sum(w, keepdims=True)
            w = w / sum_w
            age = float(np.dot(w, a))
            faceid_video_gender_age_agegroups[faceid] = [list(video_genders.keys()), face_gender, age, age_to_group(age)]
            age_face_count += 1
            age_file_count += len(video_genders)
        else:
            faceid_video_gender_age_agegroups[faceid] = [list(video_genders.keys()), face_gender, -1, UNCERTAIN]

    out_path = os.path.dirname(file_list_path)

    out_file_path = os.path.join(out_path, 'faceid_video_gender_age.txt')
    out_file = open(out_file_path, 'w')
    file_count = 0
    faceid_count = 0
    faceids = list(faceid_video_gender_age_agegroups.keys())
    faceids.sort()
    for idx, faceid in enumerate(faceids):
        video_gender_age_agegroups = faceid_video_gender_age_agegroups[faceid]
        faceid_count += 1
        videos, gender, age, agegroups = video_gender_age_agegroups
        videos.sort()
        videos_str = ','.join(videos)
        file_count += len(videos)
        out_info = '{} {} {} {} {}\n'.format(faceid, videos_str, gender, age, agegroups)
        out_file.write(out_info)
    out_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    print('Total Face Count: {}, Total File Count: {}, Age Face Count: {}, Age File Count: {}'.format(faceid_count, file_count, age_face_count, age_file_count))


if __name__ == '__main__':
    main()
